# -*- coding: utf-8 -*-
"""
:mod:`nubio.io.stockholm`
=========================

Provides read and write support for the stockholm file format.

"""
from array import array
from functools import partial
from collections import defaultdict
from tempfile import NamedTemporaryFile

from nubox.shim import izip
from nubox.char import CharArray

def stockholm_aln_parser(fh):
    """
    Parses a file or file-like ``object`` with contents in the stockholm file 
    format returns a ``tuple`` of sequences, sequence metadata and alignment
    metadata. To parse directly into a ``nubox`` or ``nubio`` data container
    use ``nubio.io.parse_aln``. 
    
    """
    seqs = None
    for line in fh:
        line = line.strip()
        if not line:
            continue
        if line.startswith('# STOCKHOLM 1.0'):
            gf = [] # per file
            gc = defaultdict(str) # per column
            ids = []
            seqs = defaultdict(str) # sequences
            gs = defaultdict(partial(defaultdict, str)) # per sequence 
            gr = defaultdict(partial(defaultdict, str)) # per sequence
        elif line.startswith('#=GF'):
            key_value = line[5:].split(None, 1)
            gf.append(key_value)
        elif line.startswith('#=GC'):
            key, value = line[5:].split(None, 1)
            gc[key] += value.strip()
        elif line.startswith('#=GS'):
            seq_id, feature, value = line[5:].split(None, 2)
            gs[seq_id][feature] += value.strip()
        elif line.startswith('#=GR'):
            seq_id, feature, value = line[5:].split(None, 2)
            gr[seq_id][feature] += value.strip()
        elif line.startswith('//'):
            # per-sequence
            gr = dict([(id, dict(feat)) for id, feat in gr.items()])
            gs = dict([(id, dict(feat)) for id, feat in gs.items()])
            seqs = [seqs[id] for id in ids]
            seq_metas = [{'stockholm':{'GS':gs.get(id, {}),
                                       'GR':gr.get(id, {}),
                                       'ID':id}} for id in ids]
            # per-alignment
            aln_meta = {}
            aln_meta['stockholm'] = {}
            aln_meta['stockholm']['GF'] = gf
            aln_meta['stockholm']['GC'] = dict(gc) # no defaultdict
            return (seqs, seq_metas, aln_meta)
        else:
            seq_id, seq = line.split(None, 1)
            if not seq_id in seqs:
                ids.append(seq_id)
            seqs[seq_id] += seq.strip()

def stockholm_aln_writer(fh, vector, feats=(), width=20):
    """
    Writes a file or file-like ``object`` in the stockholm file format from 
    the supplied "vector".
    
    Arguments:
    
      - fh(``file``) file-like ``object`` open for reading
      - vector(``nubox.char.CharArray``) a subclass of ``nubox.char.CharArray``
      - feats(``tuple``) [default: ``()``] additional "meta" keys to be exported
        as sequence features ``#=GS``.
      - width(``int``) [default: ``20``] minimum size of the id field
      
    """
    if not isinstance(vector, CharArray):
        seqs, seq_metas, aln_meta = vector
    else:
        if isinstance(vector.data, list):
            str_data = array('b', vector.data).tostring()
        else:
            str_data = vector.data.tostring()
        seqs = (str_data[start:stop] for start, stop in vector._childrenbounds())
        seq_metas = [kwargs['meta'] for kwargs in vector.childrenkwargs]
        aln_meta = vector.meta

    # write header
    lines = ['# STOCKHOLM 1.0\n']
    ID_SEQ = "%" + str(-width) + "s %s\n"
    GX_FEAT_VAL = "%s %" + str(-width + 5) + "s %s\n"
    # write GF
    gf = aln_meta.get('GF', ()) # GFs is a list
    for k, v in gf:
        lines.append("#=GF %s %s\n" % (k, v))
    # write sequences
    for i, (seq, meta) in enumerate(izip(seqs, seq_metas)):
        meta = meta.get('stockholm', {})
        id = meta.get('ID', i)
        # write sequence
        lines.append(ID_SEQ % (id, seq))
        # write GR, id exists
        for feat, value in meta.get('GR', {}).iteritems():
            feat = id + " " + feat
            lines.append(GX_FEAT_VAL % ('#=GR', feat, value))
        # write GS, if exists
        for feat, value in meta.get('GS', {}).iteritems():
            # from metadata['GS'][feat]
            feat = id + " " + feat
            lines.append(GX_FEAT_VAL % ('#=GS', feat, value))
        for i, value in enumerate(meta.get('comments', ())):
            # from metadata['comments']
            feat = id + " " + "FC%s" % i # FastaComment
            lines.append(GX_FEAT_VAL % ('#=GS', feat, value))
        for feat in feats:
            # from metadata[feat]
            value = meta.get(feat)
            if value:
                feat = id + " " + feat
                lines.append(GX_FEAT_VAL % ('#=GS', feat, value))

    # write GC
    gc = aln_meta.get('stockholm', {}).get('GC', {}) # GFs is a list
    for k, v in gc.iteritems():
        lines.append(GX_FEAT_VAL % ('#=GC', k, v))
    # write //
    lines.append('//\n')
    fh.write("".join(lines))

def stockholm_tmp(vector, **kwargs):
    """
    Writes sequence alignment to a temporary file. Returns a file handle open 
    for reading and writing. If the handle is closed or garbage collected the 
    temporary file will be deleted from the file system. Contents of the file 
    will be in the stockholm file format.
    
    Arguments:
    
        - vector(``CharArray`` or iterable). This should be a valid input for 
          ``stockholm_aln_writer``.
          
    Additional keyworded arguments are passed to ``stockholm_aln_writer``.
    
    """
    tmp_file = NamedTemporaryFile()
    stockholm_aln_writer(tmp_file, vector, **kwargs)
    tmp_file.flush()
    return tmp_file


#   Compulsory fields:
#   ------------------
#   AC   Accession number:           Accession number in form PFxxxxx (Pfam) or RFxxxxx (Rfam).
#   ID   Identification:             One word name for family.
#   DE   Definition:                 Short description of family.
#   AU   Author:                     Authors of the entry.
#   SE   Source of seed:             The source suggesting the seed members belong to one family.
#   SS   Source of structure:        The source (prediction or publication) of the consensus RNA secondary structure used by Rfam.
#   BM   Build method:               Command line used to generate the model
#   SM   Search method:              Command line used to perform the search
#   GA   Gathering method:           Search threshold to build the full alignment.
#   TC   Trusted Cutoff:             Lowest sequence score (and domain score for Pfam) of match in the full alignment.
#   NC   Noise Cutoff:               Highest sequence score (and domain score for Pfam) of match not in full alignment.
#   TP   Type:                       Type of family -- presently Family, Domain, Motif or Repeat for Pfam.
#                                                   -- a tree with roots Gene, Intron or Cis-reg for Rfam.
#   SQ   Sequence:                   Number of sequences in alignment.
#
#   Optional fields:
#   ----------------
#   DC   Database Comment:           Comment about database reference.
#   DR   Database Reference:         Reference to external database.
#   RC   Reference Comment:          Comment about literature reference.
#   RN   Reference Number:           Reference Number.
#   RM   Reference Medline:          Eight digit medline UI number.
#   RT   Reference Title:            Reference Title.
#   RA   Reference Author:           Reference Author
#   RL   Reference Location:         Journal location.
#   PI   Previous identifier:        Record of all previous ID lines.
#   KW   Keywords:                   Keywords.
#   CC   Comment:                    Comments.
#   NE   Pfam accession:            Indicates a nested domain.
#   NL   Location:                   Location of nested domains - sequence ID, start and end of insert.
#   WK   Wikipedia link:             Wikipedia page
#   CL   Clan:                       Clan accession
#   MB   Membership:                 Used for listing Clan membership
#
#   For embedding trees:
#   ----------------
#   NH  New Hampshire                A tree in New Hampshire eXtended format.
#   TN  Tree ID                      A unique identifier for the next tree.

